# @Version        : 1.0
# @Update Time    : 2024/9/15 12:49
# @File           : auth_permission.py
# @IDE            : PyCharm
# @Desc           : 权限校验类
from pydantic_core import from_json
from src.common.enum import DateScopeEnum
from src.core.exception.custom_exception import PermissionException
import json
from typing import Optional, TypeVar, List
from fastapi import Request
from jose import jwt
from sqlalchemy import ColumnElement, text, and_
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload
from src import models
from src.config import settings
from src.core.base.orm.orm import AbstractORMModel
from src.core.exception.custom_exception import AuthException
from src.schemas.user_info_schema import UserInfo
from src.log import logger
from src.services.dept_service import DeptService
from src.services.user_service import UserService
from src.db.db_redis import AsyncRedisTool

ORM = TypeVar("ORM", bound=AbstractORMModel)


class SecurityUtils:
    ALL_PERMISSION = settings.auth.ALL_PERMISSION  # 所有权限标识
    # 管理员标识
    SUPER_ADMIN = settings.auth.SUPER_ADMIN

    @staticmethod
    async def check_token(
        request: Request, token: Optional[str], redis: AsyncRedisTool
    ) -> str:
        """
        验证 JWT Token 并返回用户 ID
        """
        if not token:
            raise AuthException(msg="请先登录")
        try:
            payload = jwt.decode(
                token, settings.auth.SECRET_KEY, algorithms=[settings.auth.ALGORITHM]
            )
            session_id = payload.get("session_id")
            if not session_id:
                raise AuthException(msg="认证已失效，请您重新登录")
            user_id = await redis.get(f"{settings.cache.TOKEN_KEY}{session_id}")
            if user_id is None:
                raise AuthException(msg="认证已失效，请您重新登录")
            expire_time = await redis.client.ttl(
                f"{settings.cache.TOKEN_KEY}{session_id}"
            )
            request.state.is_refresh = -1 if expire_time < 1800 else 1
        except (jwt.JWTError, jwt.JWTClaimsError, jwt.ExpiredSignatureError):
            raise AuthException(msg="认证已失效，请您重新登录")
        return user_id

    # 从缓存或者数据库中查询用户信息
    @classmethod
    async def get_user_info(
        cls,
        request: Request,
        db: AsyncSession,
        redis: AsyncRedisTool,
        token: Optional[str],
        is_perm: bool = True,
    ) -> UserInfo:
        user_id = await cls.check_token(request, token, redis)
        user_info = await redis.get(f"{settings.cache.USER_KEY}{user_id}")
        if user_info:
            return UserInfo.model_validate(from_json(user_info, allow_partial=True))
        else:
            if is_perm:
                options = [
                    selectinload(models.SysUser.roles).subqueryload(models.SysRole.menus)
                ]
                user = await UserService(db).get_one_or_none(
                    and_(models.SysUser.id == int(user_id)), load=options
                )
                user_info = await cls.check_user_info(request, user, db, True)
                await redis.set(
                    f"{settings.cache.USER_KEY}{user_id}",
                    value=user_info.model_dump_json(),
                    ex=settings.auth.REFRESH_TOKEN_EXPIRE_MINUTES,
                )
            else:
                user = await UserService(db).get_one_or_none(
                    and_(models.SysUser.id == int(user_id))
                )
                user_info = await cls.check_user_info(request, user, db, False)
            return user_info

    @classmethod
    async def check_user_info(
        cls,
        request: Request,
        user: models.SysUser,
        db: AsyncSession,
        is_perm: bool = True,
    ) -> UserInfo:
        """
        # 校验用户信息
        :param request:
        :param user:
        :param db:
        :param is_perm:
        :return:
        """
        if not user:
            raise AuthException(msg="未认证，请您重新登陆")
        if user.status == "2":
            raise AuthException(msg="账号已禁用，请联系管理员")
        request.state.user_id = user.id
        request.state.username = user.username
        request.state.user_info = user
        try:
            request.scope["body"] = await request.body()
        except Exception as e:
            logger.error(e)
            request.scope["body"] = "获取失败"
        if not is_perm:
            return UserInfo.model_validate(user)
        data_range, dept_ids = await cls.get_user_data_range(request, user, db)
        permission = await cls.get_user_permission(user)
        uer_info = UserInfo.model_validate(user)
        uer_info.data_range = str(data_range)
        uer_info.dept_ids = dept_ids
        uer_info.permissions = permission
        return uer_info

    @classmethod
    async def get_user_permission(cls, user: models.SysUser) -> set[str]:
        """
        # 获取用户权限
        :param user:
        :return:
        """
        if user.is_super_admin():
            return {cls.ALL_PERMISSION}
        # 获取用户角色所有key值
        user_info = await AsyncRedisTool().get(f"{settings.cache.USER_KEY}{user.id}")
        user_info = json.loads(user_info)
        role_keys = user_info.get("role_keys", [])
        permissions = set()
        for role_key in role_keys:
            perm = await AsyncRedisTool().get(f"{settings.cache.ROLE_KEY}{role_key}")
            # 转成set
            perm = json.loads(perm)
            permissions.update(set(perm))
        return permissions

    @classmethod
    async def get_user_data_range(
        cls, request: Request, user: models.SysUser, db: AsyncSession
    ) -> tuple:
        """
        获取用户数据范围
        0 仅本人数据权限
        1 本部门数据权限
        2 本部门及以下数据权限
        3 自定义数据权限
        4 全部数据权限
        :param request:
        :param user:
        :param db:
        :return: 数据范围，数据范围id列表
        """
        if user.is_super == "1":
            return 4, []
        data_range = max([int(r.data_scope) for r in user.roles])
        dept_ids = set()
        if data_range == 0:
            pass
        elif data_range == 1:
            dept_ids.add(user.dept_id)
        elif data_range == 2:
            dept_ids = await DeptService(db).get_dept_and_children(dept_id=user.dept_id)
        elif data_range == 3:
            for role in user.roles:
                for dept in role.depts:
                    if dept.status == "0":
                        dept_ids.add(dept.id)
        elif data_range == 4:
            dept_ids.add("*")
        return data_range, list(dept_ids)

    @classmethod
    async def get_user_role_key(cls, request: Request, user: UserInfo) -> set:
        return set([r.role_key for r in user.roles])

    @classmethod
    async def get_date_scope(
        cls,
        request: Request,
        user: UserInfo,
        model: ORM,
        user_alias: str = "user_id",
        dept_alias: str = "dept_id",
    ) -> List[ColumnElement]:
        """
        获取数据权限sql过滤条件
        0 仅本人数据权限
        1 本部门数据权限
        2 本部门及以下数据权限
        3 自定义数据权限
        4 全部数据权限
        :param request:
        :param user:
        :param model:
        :param user_alias:
        :param dept_alias:
        :return:
        """
        where = []
        if user.is_super_admin():
            # 超级管理员拥有全部数据权限
            return where
        if user.data_range == DateScopeEnum.SELF.value:
            # 仅本人数据权限
            if hasattr(model, user_alias):
                where.append(getattr(model, user_alias) == user.id)
            else:
                where.append(text("1=-1"))
        elif (
            user.data_range == DateScopeEnum.DEPT.value
            or user.data_range == DateScopeEnum.DEPT_AND_CHILDREN.value
        ):
            # 本部门/自定义数据权限
            if hasattr(model, dept_alias):
                where.append(getattr(model, dept_alias).in_(user.dept_ids))
            else:
                where.append(text("1=-1"))
        elif user.data_range == DateScopeEnum.ALL.value:
            # 全部数据权限
            where.append(text("1=-1"))
        else:
            where.append(text("1=-1"))
        return where

    @classmethod
    def has_permi(cls, permission: str) -> bool:
        """
        校验用户是否有某权限
        :param permission:
        :return: 用户是否具有某权限
        """
        pass

    @classmethod
    def check_permi(cls, permission: str):
        """
        校验用户是否有某权限，没有则抛出异常
        :param permission:
        :return:
        """
        pass

    @classmethod
    def check_permi_and(cls, permissions: set[str], user_permissions: set[str]):
        """
        校验用户是否有某权限，必须全部拥有，没有则抛出异常
        :param permissions:需要校验的权限
        :param user_permissions:用户所有权限
        :return:
        """
        if cls.ALL_PERMISSION in user_permissions:
            return True
        # issubset: 判断permissions是否是user_permissions的子集
        if not permissions.issubset(user_permissions):
            raise PermissionException(message="用户没有权限")

    @classmethod
    def check_permi_or(cls, permissions: set[str], user_permissions: set[str]):
        """
        校验用户是否有某权限，只要有一个拥有，则通过
        :param permissions:
        :param user_permissions:
        :return:
        """
        if cls.ALL_PERMISSION in user_permissions:
            return True
        # intersection:判断permissions与user_permissions是否有交集
        if not permissions.intersection(user_permissions):
            raise PermissionException(message="用户没有权限")

    @classmethod
    def has_role(cls, role: str) -> bool:
        """
        校验用户是否有某角色
        :param role:
        :return:
        """
        pass

    @classmethod
    def check_role(cls, role: str):
        """
        校验用户是否有某角色，没有则抛出异常
        :param role:
        :return:
        """
        pass

    @classmethod
    def check_role_and(cls, roles: set, user_roles: set):
        """
        校验用户是否有某角色，必须全部拥有，没有则抛出异常
        :param roles:
        :param user_roles:
        :return:
        """
        if not roles.issubset(user_roles):
            raise PermissionException(message="用户没有权限")

    @classmethod
    def check_role_or(cls, roles: set, user_roles: set):
        """
        校验用户是否有某角色，只要有一个拥有，则通过
        :param roles:
        :param user_roles:
        :return:
        """
        if not roles.intersection(user_roles):
            raise PermissionException(message="用户没有权限")
